package com.hy.untils;

import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.completion.chat.ChatMessageRole;
import com.theokanning.openai.service.OpenAiService;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

public class OpenAIClientWrapper {
    private final OpenAiService client;

    public OpenAIClientWrapper(String apiKey, String baseUrl) {
        this.client = new OpenAiService(apiKey, Duration.parse(baseUrl));
    }

    public String call(String instruction, String prompt) {
        try {
            List<ChatMessage> messages = new ArrayList<>();
            messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(), instruction));
            messages.add(new ChatMessage(ChatMessageRole.USER.value(), prompt));

            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model("qwen1half-32b-chat")
                    .messages(messages)
                    .temperature(0.01)
                    .stream(false)
                    .presencePenalty(1.03)
                    .frequencyPenalty(1.0)
                    .topP(0.95)
                    .build();

            return client.createChatCompletion(chatCompletionRequest).getChoices().get(0).getMessage().getContent();
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }
}
